package vip.liux.backend.infrastructure.algorithm;

import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;

public class LFU<K, V> {
    // 存储键到值的映射
    private final Map<K, V> cache;
    // 存储键到频率的映射
    private final Map<K, Integer> keyFrequency;
    // 存储频率到键集合的映射
    private final Map<Integer, LinkedHashSet<K>> frequencyKeys;
    // 缓存的最大容量
    private final int capacity;
    // 当前缓存中最小的频率
    private int minFrequency;

    public LFU(int capacity) {
        this.capacity = capacity;
        this.cache = new HashMap<>();
        this.keyFrequency = new HashMap<>();
        this.frequencyKeys = new HashMap<>();
        this.minFrequency = 0;
    }

    // 获取缓存中的值
    public V get(K key) {
        if (!cache.containsKey(key)) {
            return null;
        }
        // 更新键的频率
        updateFrequency(key);
        return cache.get(key);
    }

    // 向缓存中添加键值对
    public void put(K key, V value) {
        if (capacity <= 0) {
            return;
        }
        if (cache.containsKey(key)) {
            // 如果键已经存在，更新值并更新频率
            cache.put(key, value);
            updateFrequency(key);
            return;
        }
        if (cache.size() >= capacity) {
            // 如果缓存已满，移除最不经常使用的键
            removeLeastFrequent();
        }
        // 添加新键值对
        cache.put(key, value);
        keyFrequency.put(key, 1);
        frequencyKeys.computeIfAbsent(1, k -> new LinkedHashSet<>()).add(key);
        minFrequency = 1;
    }

    // 更新键的频率
    private void updateFrequency(K key) {
        int frequency = keyFrequency.get(key); // 获取键的频率
        keyFrequency.put(key, frequency + 1); // 更新键的频率 + 1
        frequencyKeys.get(frequency).remove(key); // 从旧频率对应的键集合中移除键
        if (frequencyKeys.get(frequency).isEmpty()) {
            frequencyKeys.remove(frequency); // 如果旧频率对应的键集合为空，移除该频率
            if (minFrequency == frequency) {
                minFrequency++; // 如果移除的是最小频率，更新最小频率 + 1
            }
        }
        frequencyKeys.computeIfAbsent(frequency + 1, k -> new LinkedHashSet<>()).add(key); // 将键添加到新频率对应的键集合
    }

    // 移除最不经常使用的键
    private void removeLeastFrequent() {
        K leastFrequentKey = frequencyKeys.get(minFrequency).iterator().next(); // 获取最不经常使用的键
        frequencyKeys.get(minFrequency).remove(leastFrequentKey); // 从最小频率对应的键集合中移除键
        if (frequencyKeys.get(minFrequency).isEmpty()) {
            frequencyKeys.remove(minFrequency); // 如果最小频率对应的键集合为空，移除该频率
        }
        cache.remove(leastFrequentKey); // 从缓存中移除键值对
        keyFrequency.remove(leastFrequentKey); // 从键频率映射中移除键
    }
}